import torch
from .base_model import BaseModel
from . import networks
import numpy as np


class Blur2LambdaModel(BaseModel):
    """This is a subclass of Pix2PixModel for blur image spectralization (blur image -> spectral images cube).

    The model training requires '-dataset_model spectralization' dataset.
    It trains a pix2pix model, mapping from 1 channel to multi channels.
    By default, the spectralization dataset will automatically set '--input_nc ' and '--output_nc '.
    """

    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """Add new dataset-specific options, and rewrite default values for existing options.

        Parameters:
            parser          -- original option parser
            is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.

        Returns:
            the modified parser.

        By default, we use 'colorization' dataset for this model.
        See the original pix2pix paper (https://arxiv.org/pdf/1611.07004.pdf) and colorization results (Figure 9 in the paper)
        """
        parser.set_defaults(norm='batch', dataset_mode='spectralization', netG='TSA', netD='n_layers', n_layers_D=3)
        if is_train:
            parser.set_defaults(pool_size=0, gan_mode='vanilla')
            parser.add_argument('--lambda_L1', type=float, default=12.0, help='weight for L1 loss')
        return parser

    def __init__(self, opt):
        """Initialize the class.

        Parameters:
            opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions

        """
        # reuse the pix2pix model
        BaseModel.__init__(self, opt)

        # default stage:1
        self.stage = 1
        self.loss_R_L1 = 0
        # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses>
        self.loss_names = ['G_GAN', 'G_L1', 'D_real', 'D_fake', 'R_L1']
        # specify the images to be visualized.
        self.visual_names = ['real_A', 'fake_B_1', 'real_B_1', 'fake_B_15', 'real_B_15', 'fake_B_27', 'real_B_27']
        # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>
        if self.isTrain:
            self.model_names = ['G', 'D', 'R']
        else:  # during test time, load G and R
            self.model_names = ['G', 'R']
        # define networks (both generator and discriminator)
        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.norm,
                                      False, opt.init_type, opt.init_gain, self.gpu_ids)
        self.netR = networks.define_G(opt.output_nc, opt.output_nc, opt.ngf, 'unet_128', opt.norm,
                                      not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids)

        if self.isTrain:  # define a discriminator; conditional GANs need to take both input and output images; Therefore, #channels for D is input_nc + output_nc
            self.netD = networks.define_D(opt.input_nc + opt.output_nc, opt.ndf, opt.netD,
                                          opt.n_layers_D, opt.norm, opt.init_type, opt.init_gain, self.gpu_ids)

        if self.isTrain:
            # define loss functions
            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)
            self.criterionL1 = torch.nn.L1Loss()
            # self.criterionL1 = torch.nn.MSELoss()
            # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>.
            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)

    def set_input(self, input):
        """Unpack input data from the dataloader and perform necessary pre-processing steps.

        Parameters:
            input (dict): include the data itself and its metadata information.

        The option 'direction' can be used to swap images in domain A and domain B.
        """
        AtoB = self.opt.direction == 'AtoB'
        self.real_A = input['A' if AtoB else 'B'].to(self.device)
        self.real_B = input['B' if AtoB else 'A'].to(self.device)
        self.image_paths = input['A_paths' if AtoB else 'B_paths']

    def forward(self):
        """Run forward pass; called by both functions <optimize_parameters> and <test>."""
        if self.isTrain:
            if self.stage == 1:
                self.fake_B = self.netG(self.real_A)  # G(A)
            elif self.stage == 2:
                self.fake_B = self.netR(self.netG(self.real_A))  # R(G(A))
        else:
            self.fake_B = self.netR(self.netG(self.real_A))  # R(G(A))

    def backward_D(self):
        """Calculate GAN loss for the discriminator"""
        # Fake; stop backprop to the generator by detaching fake_B
        fake_AB = torch.cat((self.real_A, self.fake_B),
                            1)  # we use conditional GANs; we need to feed both input and output to the discriminator
        pred_fake = self.netD(fake_AB.detach())
        self.loss_D_fake = self.criterionGAN(pred_fake, False)
        # Real
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        self.loss_D_real = self.criterionGAN(pred_real, True)
        # combine loss and calculate gradients
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()

    def backward_G(self):
        """Calculate GAN and L1 loss for the generator"""
        # First, G(A) should fake the discriminator
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)
        # Second, G(A) = B
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
        # combine loss and calculate gradients
        self.loss_G = self.loss_G_GAN + self.loss_G_L1
        self.loss_G.backward()

    def backward_R(self):
        """Calculate L1 loss for the refinement"""
        self.loss_R_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1
        # combine loss and calculate gradients
        self.loss_R = self.loss_R_L1
        self.loss_R.backward()

    def optimize_parameters(self):
        self.forward()  # compute fake images: G(A)
        if self.stage == 1:
            # update D
            self.set_requires_grad(self.netD, True)  # enable backprop for D
            self.optimizer_D.zero_grad()  # set D's gradients to zero
            self.backward_D()  # calculate gradients for D
            self.optimizer_D.step()  # update D's weights
            # update G
            self.set_requires_grad(self.netD, False)  # D requires no gradients when optimizing G
            self.optimizer_G.zero_grad()  # set G's gradients to zero
            self.backward_G()  # calculate graidents for G
            self.optimizer_G.step()  # udpate G's weights
        elif self.stage == 2:
            # update R
            self.optimizer_R.zero_grad()  # set G's gradients to zero
            self.backward_R()  # calculate graidents for G
            self.optimizer_R.step()  # udpate G's weights

    def prepare_stage2(self, opt):
        self.netD.eval()  # D requires no gradients when optimizing R
        self.netG.eval()  # G requires no gradients when optimizing R

        self.optimizer_R = torch.optim.Adam(self.netR.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
        self.optimizers.pop()
        self.optimizers.pop()
        self.optimizers.append(self.optimizer_R)

        self.criterionL1 = torch.nn.L1Loss()

        self.stage = 2

        self.setup(opt)

    def strench_hw_image(self, image):
        """Strench 1-channel gray image"""
        image = torch.unsqueeze(image, 0)
        image = image.data.cpu().float().numpy()
        image_min = np.min(image)
        image_max = np.max(image)
        max_min = image_max - image_min
        stretched_img = (image - image_min) / max_min
        stretched_img = np.transpose(stretched_img.astype(np.float32), (1, 2, 0)) * 255
        stretched_img = np.repeat(stretched_img, 3, axis=2)
        return stretched_img

    def compute_visuals(self):
        """Calculate additional output images for visdom and HTML visualization"""
        self.real_B_1 = self.strench_hw_image(self.real_B[0, 0, :, :])
        self.fake_B_1 = self.strench_hw_image(self.fake_B[0, 0, :, :])
        self.real_B_15 = self.strench_hw_image(self.real_B[0, 14, :, :])
        self.fake_B_15 = self.strench_hw_image(self.fake_B[0, 14, :, :])
        self.real_B_27 = self.strench_hw_image(self.real_B[0, 26, :, :])
        self.fake_B_27 = self.strench_hw_image(self.fake_B[0, 26, :, :])

    def get_spectral_radiance_dict(self):
        """Export spectral_cube_dict for mat file """
        cube = self.fake_B[0, :, :, :]
        # adapt MATLAB CxHxW -》HxWxC
        cube = cube.permute(1, 2, 0)
        cube = cube.cpu().detach().numpy()
        bands = np.asarray(self.opt.out_channel)
        return {'cube': cube, 'bands': bands}
